# -*-coding: utf-8 -*-
import sys
import os

sys.path.insert(0, os.path.dirname(__file__))
import glob
import cv2
import numpy as np
import torch
import PIL.Image as Image
import dfsd as dfsd
import face_landmark.net.onet_landmark as onet_landmark

print("torch version:{}".format(torch.__version__))


class FaceLandmarkDetection(object):

    def __init__(self, prob_threshold=0.6, iou_threshold=0.1):
        """
        :param prob_threshold:
        """
        self.model_root = os.path.dirname(__file__)
        model_path = os.path.join(self.model_root, "data/pretrained/dsfd_vgg_0.880.pth")
        onet_path = os.path.join(self.model_root, "face_landmark/XMC2-landmark-detection.pth.tar")
        net_type = "vgg"
        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        print("use device:{}".format(self.device))
        self.face_det = dfsd.Detector(model_path=model_path,
                                      net_type=net_type,
                                      prob_threshold=prob_threshold,
                                      iou_threshold=iou_threshold,
                                      device=self.device)

        self.lmdet = onet_landmark.ONetLandmark(onet_path, device="cuda:0")

    def detect(self, rgb_image, vis=False):
        """
        :param rgb_image:
        :param vis:
        :return:
        """
        bboxes, labels, scores = self.face_det.detect_image(rgb_image, vis=False)
        if len(bboxes) == 0 or len(scores) == 0:
            return np.asarray([]), np.asarray([]), np.asarray([])
        landmarks = self.lmdet.get_image_landmarks(rgb_image, bboxes)
        if vis:
            # show bbox_score and bounding boxes
            print("bbox:\n{}\nscore:\n{}\nlandmarks:\n{}".format(bboxes, scores, landmarks))
            self.show_landmark_boxes("detect", rgb_image, landmarks, bboxes)
        return bboxes, scores, landmarks

    def detect_face_landmark(self, rgb_image, isshow=False):
        """
        :param rgb_image:
        :param isshow:
        :return:
        """
        # bboxes, scores, landmarks = self.detect(rgb_image, isshow=isshow)
        result = self.encode_dets(*self.detect(rgb_image, vis=False))
        if isshow:
            bboxes, scores, landmarks = self.decode_dets(result)
            print("bbox:\n{}\nscore:\n{}\nlandmarks:\n{}".format(bboxes, scores, landmarks))
            self.show_landmark_boxes("detect_face_landmark", rgb_image, landmarks, bboxes, scores)
        return result

    def encode_dets(self, bboxes, scores, landmarks):
        """
        :param bboxes:
        :param scores:
        :param landmarks:
        :return:
        """
        if isinstance(bboxes, np.ndarray):
            bboxes = bboxes.tolist()
        if isinstance(scores, np.ndarray):
            scores = scores.reshape(-1).tolist()
        if isinstance(landmarks, np.ndarray):
            landmarks = landmarks.tolist()
        result = []
        for i in range(len(bboxes)):
            per_result = {}
            per_result['bbox'] = {
                'left': bboxes[i][0],
                'top': bboxes[i][1],
                'right': bboxes[i][2],
                'bottom': bboxes[i][3],
                'score': scores[i]
            }
            per_result["landmark"] = []
            for land in landmarks[i]:
                l = {"x": land[0], "y": land[1]}
                per_result["landmark"].append(l)
            result.append(per_result)
        return result

    def decode_dets(self, dets):
        """
        :param dets:
        :return:
        """
        bboxes = []
        scores = []
        landmarks = []
        for item in dets:
            bbox = [item['bbox']["left"], item['bbox']["top"], item['bbox']["right"], item['bbox']["bottom"]]
            score = item['bbox']["score"]
            landmark = []
            for land in item['landmark']:
                landmark.append([land["x"], land["y"]])
            bboxes.append(bbox)
            scores.append(score)
            landmarks.append(landmark)
        bboxes = np.asarray(bboxes, dtype=np.float32)
        scores = np.asarray(scores, dtype=np.float32)
        landmarks = np.asarray(landmarks, dtype=np.float32)
        return bboxes, scores, landmarks

    @staticmethod
    def show_landmark_boxes(win_name, image, landmarks_list, boxes, scores):
        '''
        显示landmark和boxes
        :param win_name:
        :param image:
        :param landmarks_list: [[x1, y1], [x2, y2]]
        :param boxes: [[ x1, y1, x2, y2],[ x1, y1, x2, y2]]
        :return:
        '''
        point_size = 1
        point_color = (0, 0, 255)  # BGR
        thickness = 4  # 可以为 0 、4、8
        for landmarks in landmarks_list:
            for landmark in landmarks:
                # 要画的点的坐标
                point = (int(landmark[0]), int(landmark[1]))
                cv2.circle(image, point, point_size, point_color, thickness)
        FaceLandmarkDetection.show_image_boxes(image, boxes, scores)

    @staticmethod
    def show_image_boxes(image, boxes_list, scores=None):
        '''
        :param image: RGB Image
        :param boxes_list:[[ x1, y1, x2, y2],[ x1, y1, x2, y2]]
        :return:
        '''
        if scores is None:
            scores = [""] * len(boxes_list)
        for box, s in zip(boxes_list, scores):
            x1, y1, x2, y2 = box
            point1 = (int(x1), int(y1))
            point2 = (int(x2), int(y2))
            cv2.rectangle(image, point1, point2, (0, 0, 255), thickness=2)
            cv2.putText(image, str(s)[:4], point1, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), thickness=2)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        cv2.imshow("Det", image)
        cv2.waitKey(0)


def demo(image_file):
    """
    :param image_file: image file path or directory
    :return:
    """
    det = FaceLandmarkDetection()
    if os.path.isdir(image_file):
        image_list = glob.glob(os.path.join(image_file, "*"))
    else:
        image_list = [image_file]
    for image_path in image_list:
        print(image_path)
        image = cv2.imread(image_path)
        if image.size == 0:
            continue
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        result = det.detect_face_landmark(image, isshow=True)
        print(result)


if __name__ == "__main__":
    image_path = "./data/test_image"
    demo(image_path)
